Skip to content

feat(timm): enable timm image-classification models via library routing#790

Merged
vortex-captain merged 10 commits into
mainfrom
reny/timm
Jun 2, 2026
Merged

feat(timm): enable timm image-classification models via library routing#790
vortex-captain merged 10 commits into
mainfrom
reny/timm

Conversation

@vortex-captain

@vortex-captain vortex-captain commented May 29, 2026

Copy link
Copy Markdown
Contributor

Summary

timm checkpoints load through transformers'' generic TimmWrapper (model_type="timm_wrapper") and previously failed in every winml command with "Cannot detect task: config has no ''architectures'' field". Two gaps:

  1. Task/class detection — timm repos load as TimmWrapperConfig with architectures=None, so auto-detection could not resolve a task or class.
  2. OnnxConfig location — Optimum registers timm''s config (TimmDefaultOnnxConfig) only under library_name="timm", but every winml lookup defaults to transformers.

timm_wrapper is transformers'' generic bridge for the whole timm library — not a model architecture — so it is resolved at the shared resolution layer, not as a per-model config. Only the library is recorded; the task is derived from Optimum.

Changes (no models/hf/ entry)

  • loader/task.pyWRAPPED_LIBRARY_MODEL_TYPES (model_type -> optimum_library) + resolve_optimum_library(). When a config has no architectures, _detect_task_and_class_from_config derives the task from Optimum''s task list for the library (get_supported_tasks("timm_wrapper", "timm") -> ["image-classification"]) and the class from get_model_class_for_task (generic AutoModelForImageClassification, which transformers dispatches to TimmWrapper at load). The task is not hardcoded; the branch imports optimum.exporters.onnx.model_configs first to populate Optimum''s registry (scoped so normal model loading never pays for it).
  • export/io.py_get_onnx_config routes the library via resolve_optimum_library, so timm_wrapper resolves Optimum''s TimmDefaultOnnxConfig from every call site (config/build/export/inspect) with no --library flag.
  • commands/inspect.py + inspect/resolver.py — route both the CLI inspect path and the public inspect_model path the same way: library routing for the OnnxConfig lookup, plus wrapped-library task detection so the task is not mislabeled.
  • Tests: resolve_optimum_library + wrapped-library architectures fallback with task derivation (loader); timm library routing for resolve_io_specs / _get_onnx_config (export); public inspect path detect_task / resolve_exporter for timm (inspect).

Validation

Functional (end-to-end) on a timm image-classification model:

Command Before After
winml config exit 2 — no ''architectures'' field task=image-classification, 1 input
winml export exit 2 — same model.onnx (pixel_values to logits)
winml inspect exit 1 — same AutoModelForImageClassification + TimmDefaultOnnxConfig, full I/O table

config -> export -> optimize -> model.onnx validated end-to-end for multiple timm CNN classifiers. Also resolves on a timm ViT backbone (num_labels=0) -> task=image-classification, matching Optimum''s own infer_task_from_model, so it generalizes across timm architectures (CNN + ViT).

No impact on existing models — scanned all 439 entries / 401 unique models in scripts/e2e_eval/testsets/models_all.json: 0 are timm_wrapper (by JSON metadata and by loaded config; 330 loadable). Since timm_wrapper is the only trigger of the new branch, no existing model changes behavior. (71 fail to load a config — custom/GGUF/tabular types that fail at AutoConfig regardless; 7 have empty architectures but are not timm — a pre-existing "Cannot detect task", identical before and after the PR.)

No overhead for normal (non-timm) modelswinml config on a standard non-timm model: this branch vs base, min ~12.6s vs ~12.5s (within run-to-run noise). Non-timm configs have architectures, so they skip the new branch; the only added cost is one dict lookup.

Unit teststests/unit/loader + tests/unit/export + tests/unit/inspect: green.

🤖 Generated with Claude Code

Comment thread tests/unit/export/test_timm_wrapper_onnx_config.py Fixed
@vortex-captain vortex-captain changed the title feat(timm): enable timm image-classification models across CLI commands feat(timm): enable timm image-classification models via library routing May 29, 2026
Comment thread tests/unit/export/test_timm_library_routing.py Dismissed
timm checkpoints (e.g. timm/mobilenetv3_small_100.lamb_in1k) load through
transformers'' generic TimmWrapper (model_type="timm_wrapper") and previously
failed in every winml command with "Cannot detect task: config has no
''architectures'' field":

1. timm repo configs load as TimmWrapperConfig with architectures=None, so
   task/model-class auto-detection could not resolve anything.
2. Optimum registers timm''s OnnxConfig (TimmDefaultOnnxConfig) only under
   library_name="timm"; every winml lookup defaults to "transformers".

timm_wrapper is transformers'' generic bridge for the whole timm library, not a
model architecture, so it is handled at the shared resolution layer rather than
as a per-model config. Only the library is recorded; the task is derived.

- loader/task.py: add WRAPPED_LIBRARY_MODEL_TYPES (model_type -> optimum_library)
  + resolve_optimum_library(). When a config has no `architectures`,
  _detect_task_and_class_from_config derives the task from Optimum''s task list
  for that library (get_supported_tasks) and the class from
  get_model_class_for_task (a generic AutoModelFor* that transformers dispatches
  to TimmWrapper at load).
- export/io.py: _get_onnx_config routes the library via resolve_optimum_library,
  so timm_wrapper resolves Optimum''s TimmDefaultOnnxConfig from every call site
  (config/build/export/inspect) without a --library flag.
- commands/inspect.py: route the display-label lookup the same way.

Tests: resolve_optimum_library + wrapped-library architectures fallback with
task derivation (loader); timm library routing for resolve_io_specs /
_get_onnx_config (export). Verified config/export/inspect on
timm/mobilenetv3_small_100.lamb_in1k.

@vortex-captain vortex-captain left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review: feat(timm) — enable timm image-classification via library routing

I reviewed this fresh (treating nothing as known-good) and verified the core Optimum taxonomy facts by read-only introspection of the installed optimum 2.1.0 / optimum-onnx 0.1.0 / transformers 4.57.6 stack:

  • get_supported_tasks_for_model_type("timm_wrapper", exporter="onnx", library_name="timm")["image-classification"] (exactly one task).
  • The same lookup under library_name="transformers" raises KeyError("timm_wrapper is not supported yet for transformers"), so get_supported_tasks (which swallows exceptions → []) correctly does not false-trigger the new branch for non-timm/default paths.
  • The scoped import optimum.exporters.onnx.model_configs is genuinely load-bearing in the loader branch: without it, get_supported_tasks("timm_wrapper","timm") raises KeyError('default-timm-config') → returns [] → branch skipped. Good that the import is placed there.
  • TimmWrapperConfig(num_labels=10) really does have model_type="timm_wrapper" and architectures=None, so the test fixtures are realistic.
  • Both timm_wrapper and default-timm-config resolve to TimmDefaultOnnxConfig under library="timm".
  • export → loader import direction already exists on main (export/config.py, export/htp/exporter.py), so the new from ..loader.task import … introduces no new layering violation.

Verdict: the mechanism is sound and the centralization in _get_onnx_config is a clean choice. Below are correctness edge-cases, one consistency gap, and several minor/convention items. Comment-only — no blocking objection.

Notable item not on a changed line (so flagged here)

src/winml/modelkit/inspect/resolver.py:346 — the parallel inspect path resolve_exporter() still hardcodes library_name="transformers" in its TasksManager.get_exporter_config_constructor(...) call. This function is exported in the inspect package public API (inspect/__init__.py __all__inspect_modelresolve_exporter) but is not the path the winml inspect CLI uses (the CLI calls commands/inspect.py::_inspect_model_v2, which you did update). Still, for a timm model the legacy inspect_model() would (a) mislabel the OnnxConfig (the constructor lookup raises KeyError under transformers and is swallowed) and (b) resolver.py::detect_task would fall through its except ValueError to HF_TASK_DEFAULTS and return next-sentence-prediction as the task. Since resolve_io_specs/_discover_io_attrs_from_onnx_config now self-route via _get_onnx_config, this leaves that public path half-routed. Either route it the same way for consistency, or note that inspect_model()/resolve_exporter are legacy and slated for removal.

Comment thread src/winml/modelkit/loader/task.py
Comment thread src/winml/modelkit/loader/task.py Outdated
Comment thread src/winml/modelkit/loader/task.py Outdated
Comment thread src/winml/modelkit/loader/task.py
Comment thread src/winml/modelkit/loader/task.py
Comment thread tests/unit/loader/test_detect_task_and_class.py
Comment thread src/winml/modelkit/export/io.py
Comment thread src/winml/modelkit/commands/inspect.py
Yi Ren and others added 4 commits May 29, 2026 17:07
…ries

The CLI inspect path (_inspect_model_v2) already routes timm via
resolve_optimum_library, but the parallel public path
(inspect.inspect_model -> resolver.detect_task / resolve_exporter) still
hardcoded library_name="transformers" and mislabeled the task via the
HF_TASK_DEFAULTS fallback for timm_wrapper configs (no architectures).

- resolver.detect_task: reuse the loader wrapped-library resolution to derive
  the real task (image-classification) instead of mislabeling.
- resolver.resolve_exporter: route the OnnxConfig lookup via
  resolve_optimum_library so it resolves Optimum''s TimmDefaultOnnxConfig.

Adds tests/unit/inspect/test_resolver_timm.py covering both. Addresses PR review #1.
, #5)

- Correct the resolve_optimum_library docstring: an explicit
  `--library transformers` is still rerouted for wrapped types (the predicate is
  library_name == "transformers"); only an explicit non-transformers library is
  returned unchanged.
- Export resolve_optimum_library via loader/__init__.py (__all__ + module
  Public API docstring) since it is consumed cross-package; import it from
  ..loader in export/io.py and commands/inspect.py.
…iew #6)

WRAPPED_LIBRARY_MODEL_TYPES.get(...) already returns None for a missing/None
key, so the `or ""` only fabricated a meaningless empty model_type. Compute
model_type once and use it consistently (guarded so the lookup stays
type-safe), removing the direct config.model_type dereferences in the branch.
DingmaomaoBJTU

This comment was marked as duplicate.

DingmaomaoBJTU

This comment was marked as duplicate.

DingmaomaoBJTU

This comment was marked as duplicate.

DingmaomaoBJTU

This comment was marked as duplicate.

DingmaomaoBJTU

This comment was marked as duplicate.

DingmaomaoBJTU

This comment was marked as duplicate.

DingmaomaoBJTU

This comment was marked as duplicate.

DingmaomaoBJTU

This comment was marked as duplicate.

DingmaomaoBJTU

This comment was marked as duplicate.

DingmaomaoBJTU

This comment was marked as duplicate.

DingmaomaoBJTU

This comment was marked as duplicate.

DingmaomaoBJTU

This comment was marked as duplicate.

DingmaomaoBJTU

This comment was marked as duplicate.

DingmaomaoBJTU

This comment was marked as duplicate.

DingmaomaoBJTU

This comment was marked as duplicate.

DingmaomaoBJTU

This comment was marked as duplicate.

DingmaomaoBJTU

This comment was marked as duplicate.

DingmaomaoBJTU

This comment was marked as duplicate.

Comment thread src/winml/modelkit/loader/task.py
Comment thread src/winml/modelkit/loader/task.py
DingmaomaoBJTU

This comment was marked as duplicate.

Comment thread src/winml/modelkit/loader/task.py
Comment thread src/winml/modelkit/loader/task.py
DingmaomaoBJTU

This comment was marked as duplicate.

DingmaomaoBJTU

This comment was marked as duplicate.

@vortex-captain vortex-captain marked this pull request as ready for review June 1, 2026 05:46
@vortex-captain vortex-captain requested a review from a team as a code owner June 1, 2026 05:46
Comment thread src/winml/modelkit/loader/task.py
Comment thread src/winml/modelkit/inspect/resolver.py Outdated
Comment thread src/winml/modelkit/commands/inspect.py
vortex-captain and others added 2 commits June 1, 2026 17:08
…ource constant

Review follow-ups:
- loader/task.py: comment why supported[0] is the correct default (a wrapped
  library exposes a single ONNX export task today -- timm => image-classification),
  and that it needs revisiting if one ever exposes multiple tasks.
- inspect/resolver.py: replace the "wrapped-library" detect_task source string
  with a named WRAPPED_LIBRARY_SOURCE constant.
@vortex-captain vortex-captain requested a review from zhenchaoni June 1, 2026 09:37
DingmaomaoBJTU

This comment was marked as duplicate.

DingmaomaoBJTU

This comment was marked as duplicate.

DingmaomaoBJTU

This comment was marked as duplicate.

DingmaomaoBJTU

This comment was marked as duplicate.

DingmaomaoBJTU

This comment was marked as duplicate.

DingmaomaoBJTU

This comment was marked as duplicate.

Comment thread src/winml/modelkit/loader/task.py
Comment thread src/winml/modelkit/loader/task.py
…rt task

_detect_task_and_class_from_config still defaults to supported[0] for a
wrapped-library model with no `architectures`, but now logs a warning naming the
supported tasks when there is more than one, so the arbitrary default is visible
(pass --task to choose another). timm exposes a single task today, so the normal
path is unchanged.
@vortex-captain vortex-captain enabled auto-merge (squash) June 2, 2026 04:06
@vortex-captain vortex-captain merged commit eab2c48 into main Jun 2, 2026
9 checks passed
@vortex-captain vortex-captain deleted the reny/timm branch June 2, 2026 05:08
DingmaomaoBJTU pushed a commit that referenced this pull request Jun 5, 2026
…ng (#790)

## Summary

timm checkpoints load through transformers'' generic `TimmWrapper`
(`model_type="timm_wrapper"`) and previously failed in **every** `winml`
command with *"Cannot detect task: config has no ''architectures''
field"*. Two gaps:

1. **Task/class detection** — timm repos load as `TimmWrapperConfig`
with `architectures=None`, so auto-detection could not resolve a task or
class.
2. **OnnxConfig location** — Optimum registers timm''s config
(`TimmDefaultOnnxConfig`) only under `library_name="timm"`, but every
`winml` lookup defaults to `transformers`.

`timm_wrapper` is transformers'' generic bridge for the whole timm
library — not a model architecture — so it is resolved at the **shared
resolution layer**, not as a per-model config. Only the library is
recorded; the task is derived from Optimum.

## Changes (no `models/hf/` entry)

- **`loader/task.py`** — `WRAPPED_LIBRARY_MODEL_TYPES` (`model_type ->
optimum_library`) + `resolve_optimum_library()`. When a config has no
`architectures`, `_detect_task_and_class_from_config` derives the task
from Optimum''s task list for the library
(`get_supported_tasks("timm_wrapper", "timm")` ->
`["image-classification"]`) and the class from
`get_model_class_for_task` (generic `AutoModelForImageClassification`,
which transformers dispatches to `TimmWrapper` at load). The task is not
hardcoded; the branch imports `optimum.exporters.onnx.model_configs`
first to populate Optimum''s registry (scoped so normal model loading
never pays for it).
- **`export/io.py`** — `_get_onnx_config` routes the library via
`resolve_optimum_library`, so `timm_wrapper` resolves Optimum''s
`TimmDefaultOnnxConfig` from every call site
(config/build/export/inspect) with no `--library` flag.
- **`commands/inspect.py`** + **`inspect/resolver.py`** — route both the
CLI inspect path and the public `inspect_model` path the same way:
library routing for the OnnxConfig lookup, plus wrapped-library task
detection so the task is not mislabeled.
- Tests: `resolve_optimum_library` + wrapped-library architectures
fallback with task derivation (loader); timm library routing for
`resolve_io_specs` / `_get_onnx_config` (export); public inspect path
`detect_task` / `resolve_exporter` for timm (inspect).

## Validation

**Functional (end-to-end)** on a timm image-classification model:

| Command | Before | After |
|---|---|---|
| `winml config` | exit 2 — *no ''architectures'' field* |
task=image-classification, 1 input |
| `winml export` | exit 2 — same | `model.onnx` (pixel_values to logits)
|
| `winml inspect` | exit 1 — same | `AutoModelForImageClassification` +
`TimmDefaultOnnxConfig`, full I/O table |

`config` -> `export` -> `optimize` -> `model.onnx` validated end-to-end
for multiple timm CNN classifiers. Also resolves on a timm ViT backbone
(`num_labels=0`) -> task=image-classification, matching Optimum''s own
`infer_task_from_model`, so it generalizes across timm architectures
(CNN + ViT).

**No impact on existing models** — scanned all 439 entries / 401 unique
models in `scripts/e2e_eval/testsets/models_all.json`: **0** are
`timm_wrapper` (by JSON metadata and by loaded config; 330 loadable).
Since `timm_wrapper` is the only trigger of the new branch, no existing
model changes behavior. (71 fail to load a config — custom/GGUF/tabular
types that fail at `AutoConfig` regardless; 7 have empty `architectures`
but are not timm — a pre-existing "Cannot detect task", identical before
and after the PR.)

**No overhead for normal (non-timm) models** — `winml config` on a
standard non-timm model: this branch vs base, min ~12.6s vs ~12.5s
(within run-to-run noise). Non-timm configs have `architectures`, so
they skip the new branch; the only added cost is one dict lookup.

**Unit tests** — `tests/unit/loader` + `tests/unit/export` +
`tests/unit/inspect`: green.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Co-authored-by: Yi Ren <reny@microsoft.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants